import os.path as osp
import os
import pickle
import time

from tqdm import tqdm
import matplotlib.pyplot as plt

enabled = True
BASE_PATH = '__custom_pycache__'
FIGURE_PATH = '__figure__'
base_path = BASE_PATH


class CustomizedCachePath:
    def __init__(self, path):
        global base_path
        base_path = path

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        global base_path
        base_path = BASE_PATH


def to_cache_path(cache_id: str):
    return f'./{base_path}/{cache_id}.pkl'


def register(cache_id, data, overwrite):
    if not enabled:
        return
    if not osp.exists(base_path):
        os.mkdir(base_path)
    path = to_cache_path(cache_id)
    if not osp.exists(path) or overwrite:
        with open(path, 'wb') as fout:
            pickle.dump(data, fout)


def check(cache_id: str) -> bool:
    if not enabled:
        return False
    return osp.exists(to_cache_path(cache_id))


def load(cache_id: str, extra_check=True):
    if not enabled or not extra_check:
        return False, None
    if not osp.exists(to_cache_path(cache_id)):
        return False, None
    with open(to_cache_path(cache_id), 'rb') as fin:
        data = pickle.load(fin)
        return True, data


def save_figure(figure_id: str, *args, **kwargs):
    if not osp.exists(FIGURE_PATH):
        os.mkdir(FIGURE_PATH)
    plt.savefig(f'{FIGURE_PATH}/{figure_id}', *args, **kwargs)


if __name__ == '__main__':
    # for i in tqdm(range(100), desc='Processing', leave=True):
    #     time.sleep(0.04)
    with CustomizedCachePath('asd'):
        print(base_path)

    print(base_path)
    register(123456, [1, 2, 3, 4], True)
